package io.gatling.extensions.kafka.builder

import com.typesafe.scalalogging.StrictLogging
import io.gatling.core.action.Action
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.session.Expression
import io.gatling.core.structure.ScenarioContext
import io.gatling.extensions.kafka.action.KafkaProductAction
import io.gatling.extensions.kafka.protocol.KafkaProducerProtocol

case class KafkaProductActionBuilder(name: String) extends ActionBuilder with StrictLogging {

  private var topic: Expression[String] = _
  private var key: Option[Expression[String]] = _
  private var value: Expression[String] = _

  def topic(topic: Expression[String]): KafkaProductActionBuilder = {
    this.topic = topic
    this
  }

  def send(key: Option[Expression[String]], value: Expression[String]): KafkaProductActionBuilder = {
    this.key = key
    this.value = value
    this
  }

  override def build(ctx: ScenarioContext, next: Action): Action = {
    import ctx._
    KafkaProductAction(
      name,
      topic,
      key,
      value,
      protocolComponentsRegistry.components(KafkaProducerProtocol.kafkaProtocolKey),
      coreComponents,
      throttled,
      next
    )
  }
}
